import onnx
from onnx_tf.backend import prepare
import tensorflow as tf

def convert_tf_to_tflite(tf_saved_model_dir="tf_model", tflite_filename="temperature_humidity_mlp.tflite"):
    converter = tf.lite.TFLiteConverter.from_saved_model(tf_saved_model_dir)
    tflite_model = converter.convert()

    with open(tflite_filename, "wb") as f:
        f.write(tflite_model)

    print(f"TFLite模型已保存为 {tflite_filename}")


def convert_onnx_to_tf(onnx_filename="model.onnx", tf_saved_model_dir="tf_model"):
    model = onnx.load(onnx_filename)
    tf_rep = prepare(model)
    tf_rep.export_graph(tf_saved_model_dir)
    print(f"TensorFlow模型已保存到 {tf_saved_model_dir}")

# 调用转换函数
convert_onnx_to_tf(onnx_filename='model.onnx', tf_saved_model_dir='tf_model')

# 调用转换函数
convert_tf_to_tflite()